import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import datetime
import akshare as ak
import talib
from matplotlib.gridspec import GridSpec
import mplfinance as mpf

def calculate_macd(data, fast_period=6, slow_period=13, signal_period=5):
    """Calculate MACD"""
    close = data['close']
    ema_fast = talib.EMA(close, timeperiod=fast_period)
    ema_slow = talib.EMA(close, timeperiod=slow_period)
    dif = ema_fast - ema_slow
    dea = talib.EMA(dif, timeperiod=signal_period)
    macd = (dif - dea) * 2  # MACD histogram
    return dif, dea, macd

def calculate_kdj(data, n=81, m1=3, m2=3):
    """Calculate KDJ"""
    low_list = data['low'].rolling(window=n).min()
    high_list = data['high'].rolling(window=n).max()
    rsv = (data['close'] - low_list) / (high_list - low_list) * 100
    
    k = pd.Series(0.0, index=data.index)
    d = pd.Series(0.0, index=data.index)
    j = pd.Series(0.0, index=data.index)
    
    for i in range(n, len(data)):
        if i == n:
            k[i] = 50
            d[i] = 50
            j[i] = 50
        else:
            k[i] = (1-1/m1) * k[i-1] + 1/m1 * rsv[i]
            d[i] = (1-1/m2) * d[i-1] + 1/m2 * k[i]
            j[i] = 3 * k[i] - 2 * d[i]
    
    return k, d, j

def get_entry_signals(data):
    """Get entry signals based on MACD and KDJ"""
    signals = pd.Series(index=data.index, data=False)
    
    # MACD DIF crosses above DEA
    macd_cross = (data['dif'] > data['dea']) & (data['dif'].shift(1) <= data['dea'].shift(1))
    
    # KDJ J crosses above K or D
    kdj_cross_k = (data['j'] > data['k']) & (data['j'].shift(1) <= data['k'].shift(1))
    kdj_cross_d = (data['j'] > data['d']) & (data['j'].shift(1) <= data['d'].shift(1))
    
    # Entry signal
    signals = macd_cross & (kdj_cross_k | kdj_cross_d)
    
    return signals

def get_exit_signals(data):
    """Get exit signals based on KDJ J > 100"""
    signals = (data['j'] > 100)
    return signals

def backtest(data, initial_capital=100000):
    """Perform backtest on the given strategy"""
    
    # Initialize portfolio metrics
    positions = pd.Series(0, index=data.index)
    cash = pd.Series(initial_capital, index=data.index)
    portfolio_value = pd.Series(initial_capital, index=data.index)
    returns = pd.Series(0.0, index=data.index)
    trades = []
    
    # Apply trading signals
    current_position = 0
    entry_price = 0
    
    for i in range(1, len(data)):
        if i == 0:
            positions[i] = 0
            cash[i] = initial_capital
            portfolio_value[i] = initial_capital
            continue
        
        positions[i] = positions[i-1]
        cash[i] = cash[i-1]
        
        # Entry signal
        if data['entry_signal'][i] and current_position == 0:
            # Calculate max shares to buy with 95% of cash
            max_shares = int((cash[i] * 0.95) / data['close'][i])
            if max_shares > 0:
                current_position = max_shares
                positions[i] = current_position
                entry_price = data['close'][i]
                cash[i] -= entry_price * current_position
                trades.append({
                    'type': 'BUY',
                    'time': data.index[i],
                    'price': entry_price,
                    'shares': current_position,
                    'value': entry_price * current_position
                })
        
        # Exit signal
        elif data['exit_signal'][i] and current_position > 0:
            exit_price = data['close'][i]
            cash[i] += exit_price * current_position
            trades.append({
                'type': 'SELL',
                'time': data.index[i],
                'price': exit_price,
                'shares': current_position,
                'value': exit_price * current_position,
                'profit': (exit_price - entry_price) * current_position
            })
            current_position = 0
            positions[i] = 0
        
        # Calculate portfolio value
        portfolio_value[i] = cash[i] + positions[i] * data['close'][i]
        
        # Calculate returns
        returns[i] = (portfolio_value[i] / portfolio_value[i-1]) - 1
    
    # Calculate metrics
    data['positions'] = positions
    data['cash'] = cash
    data['portfolio_value'] = portfolio_value
    data['returns'] = returns
    data['cumulative_returns'] = (1 + returns).cumprod() - 1
    
    return data, trades

def plot_backtest_results(data, trades):
    """Plot backtest results"""
    # Create figure with multiple subplots
    fig = plt.figure(figsize=(16, 14))
    gs = GridSpec(4, 1, figure=fig, height_ratios=[3, 1, 1, 1])
    
    # Price subplot with buy/sell signals
    ax1 = fig.add_subplot(gs[0])
    ax1.plot(data.index, data['close'], color='black', lw=1.5)
    ax1.set_title('510310-沪深300ETF Price with Trade Signals')
    ax1.set_ylabel('Price')
    ax1.grid(True)
    
    # Plot buy signals
    buy_signals = [trade for trade in trades if trade['type'] == 'BUY']
    if buy_signals:
        buy_times = [trade['time'] for trade in buy_signals]
        buy_prices = [trade['price'] for trade in buy_signals]
        ax1.scatter(buy_times, buy_prices, marker='^', color='g', s=100, label='Buy')
    
    # Plot sell signals
    sell_signals = [trade for trade in trades if trade['type'] == 'SELL']
    if sell_signals:
        sell_times = [trade['time'] for trade in sell_signals]
        sell_prices = [trade['price'] for trade in sell_signals]
        ax1.scatter(sell_times, sell_prices, marker='v', color='r', s=100, label='Sell')
    
    ax1.legend()
    
    # MACD subplot
    ax2 = fig.add_subplot(gs[1], sharex=ax1)
    ax2.plot(data.index, data['dif'], color='blue', lw=1.5, label='DIF (Fast)')
    ax2.plot(data.index, data['dea'], color='orange', lw=1.5, label='DEA (Slow)')
    ax2.bar(data.index, data['macd'], color='green', alpha=0.5)
    ax2.set_ylabel('MACD')
    ax2.grid(True)
    ax2.legend()
    
    # KDJ subplot
    ax3 = fig.add_subplot(gs[2], sharex=ax1)
    ax3.plot(data.index, data['k'], color='blue', lw=1.5, label='K')
    ax3.plot(data.index, data['d'], color='orange', lw=1.5, label='D')
    ax3.plot(data.index, data['j'], color='purple', lw=1.5, label='J')
    ax3.axhline(y=100, color='r', linestyle='--', alpha=0.5, label='J=100')
    ax3.set_ylabel('KDJ')
    ax3.grid(True)
    ax3.legend()
    
    # Portfolio value
    ax4 = fig.add_subplot(gs[3], sharex=ax1)
    ax4.plot(data.index, data['portfolio_value'], color='green', lw=2)
    ax4.set_ylabel('Portfolio Value')
    ax4.set_xlabel('Date')
    ax4.grid(True)
    
    # Set x-axis format
    for ax in [ax1, ax2, ax3, ax4]:
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M'))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
    
    plt.tight_layout()
    plt.show()

def print_performance_metrics(data, trades):
    """Print performance metrics"""
    initial_capital = data['cash'][0]
    final_value = data['portfolio_value'].iloc[-1]
    total_return = (final_value / initial_capital - 1) * 100
    
    # Calculate daily returns
    daily_returns = data['returns'].resample('D').sum()
    
    # Annualized return
    days = (data.index[-1] - data.index[0]).days
    if days > 0:
        annualized_return = ((1 + total_return/100) ** (365/days) - 1) * 100
    else:
        annualized_return = 0
    
    # Max drawdown
    cum_returns = data['cumulative_returns']
    running_max = cum_returns.cummax()
    drawdown = (cum_returns - running_max) * 100
    max_drawdown = drawdown.min()
    
    # Trade statistics
    num_trades = len([t for t in trades if t['type'] == 'BUY'])
    profitable_trades = len([t for t in trades if t['type'] == 'SELL' and t.get('profit', 0) > 0])
    if num_trades > 0:
        win_rate = (profitable_trades / num_trades) * 100
    else:
        win_rate = 0
    
    # Profit/loss
    total_profit = sum([t.get('profit', 0) for t in trades if t['type'] == 'SELL' and t.get('profit', 0) > 0])
    total_loss = sum([t.get('profit', 0) for t in trades if t['type'] == 'SELL' and t.get('profit', 0) < 0])
    
    print("\n===== PERFORMANCE METRICS =====")
    print(f"Initial Capital: ¥{initial_capital:,.2f}")
    print(f"Final Portfolio Value: ¥{final_value:,.2f}")
    print(f"Total Return: {total_return:.2f}%")
    print(f"Annualized Return: {annualized_return:.2f}%")
    print(f"Max Drawdown: {max_drawdown:.2f}%")
    print(f"Number of Trades: {num_trades}")
    print(f"Profitable Trades: {profitable_trades}")
    print(f"Win Rate: {win_rate:.2f}%")
    if num_trades > 0:
        print(f"Average Profit per Trade: ¥{(total_profit + total_loss) / num_trades:,.2f}")
    print("=============================")

def run_backtest():
    """Main function to run the backtest"""
    print("Starting backtest for 510310-沪深300ETF...")
    
    # Get current date and 30 days ago for 1-month backtesting
    end_date = datetime.datetime.now()
    start_date = end_date - datetime.timedelta(days=30)
    
    print(f"Fetching data from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}...")
    
    try:
        # Get 1-minute data for 510310 using akshare
        stock_code = "sh510310"  # 510310 in Shanghai exchange
        data = ak.stock_zh_a_minute(symbol=stock_code, period="1", adjust="qfq")
        
        # Rename columns and process data
        data.columns = ['datetime', 'open', 'high', 'low', 'close', 'volume', 'amount']
        data['datetime'] = pd.to_datetime(data['datetime'])
        data.set_index('datetime', inplace=True)
        
        # Filter for the last month
        data = data.loc[start_date:end_date]
        
        print(f"Data loaded successfully. Got {len(data)} data points.")
        
        if len(data) == 0:
            print("No data found for the specified period.")
            return
        
        print("Calculating indicators...")
        # Calculate MACD with specified parameters (6, 13, 5)
        data['dif'], data['dea'], data['macd'] = calculate_macd(data, fast_period=6, slow_period=13, signal_period=5)
        
        # Calculate KDJ with specified parameters (81, 3, 3)
        data['k'], data['d'], data['j'] = calculate_kdj(data, n=81, m1=3, m2=3)
        
        # Get entry and exit signals
        print("Generating trading signals...")
        data['entry_signal'] = get_entry_signals(data)
        data['exit_signal'] = get_exit_signals(data)
        
        # Run backtest
        print("Running backtest...")
        backtest_results, trades = backtest(data)
        
        # Print results and plot
        print_performance_metrics(backtest_results, trades)
        plot_backtest_results(backtest_results, trades)
        
    except Exception as e:
        print(f"Error during backtesting: {e}")

if __name__ == "__main__":
    run_backtest()
